use std::{collections::HashMap, str::FromStr};

use serde::Deserialize;

use crate::errors::OapiError;

/// The streaming and non-streaming completion response shares the same json structure.
#[derive(Debug, Deserialize, Clone)]
pub struct Completion {
    /// A unique identifier for the completion.
    pub id: String,
    /// The list of completion choices the model generated for the input prompt.
    pub choices: Vec<CompletionChoice>,
    /// The Unix timestamp (in seconds) of when the completion was created.
    pub created: usize,
    /// The model used for completion.
    pub model: String,
    /// The object type, which is always "text_completion"
    pub object: String,
    /// This fingerprint represents the backend configuration that the model runs with.
    ///
    /// Can be used in conjunction with the `seed` request parameter to understand when
    /// backend changes have been made that might impact determinism.
    pub system_fingerprint: Option<String>,
    /// Usage statistics for the completion request.
    pub usage: Option<CompletionUsage>,
}

#[derive(Debug, Deserialize, Clone)]
pub struct Logprobs {
    /// The offset into the generated text for each token.
    pub text_offset: Option<Vec<usize>>,
    /// The log probability of each token in the generated text.
    pub token_logprobs: Option<Vec<f32>>,
    /// The tokens generated by the model.
    pub tokens: Option<Vec<String>>,
    /// The top log probabilities for each token position.
    pub top_logprobs: Option<Vec<HashMap<String, f32>>>,
}

#[derive(Debug, Deserialize, Clone)]
pub struct CompletionChoice {
    /// The reason the model stopped generating tokens.
    pub finish_reason: Option<String>,
    /// The index of this choice in the array of choices.
    pub index: usize,
    /// The log probabilities for each token in the generated text.
    pub logprobs: Option<Logprobs>,
    /// The generated text.
    pub text: String,
}

#[derive(Debug, Deserialize, Clone)]
pub struct CompletionTokensDetails {
    /// When using Predicted Outputs, the number of tokens in the prediction that
    /// appeared in the completion.
    pub accepted_prediction_tokens: Option<usize>,

    /// Audio input tokens generated by the model.
    pub audio_tokens: Option<usize>,

    /// Tokens generated by the model for reasoning.
    pub reasoning_tokens: Option<usize>,

    /// When using Predicted Outputs, the number of tokens in the prediction that did
    /// not appear in the completion. However, like reasoning tokens, these tokens are
    /// still counted in the total completion tokens for purposes of billing, output,
    /// and context window limits.
    pub rejected_prediction_tokens: Option<usize>,
}

#[derive(Debug, Deserialize, Clone)]
pub struct PromptTokensDetails {
    /// Audio input tokens present in the prompt.
    pub audio_tokens: Option<usize>,

    /// Cached tokens present in the prompt.
    pub cached_tokens: Option<usize>,
}

#[derive(Debug, Deserialize, Clone)]
pub struct CompletionUsage {
    /// Number of tokens in the generated completion.
    pub completion_tokens: usize,

    /// Number of tokens in the prompt.
    pub prompt_tokens: usize,

    /// Total number of tokens used in the request (prompt + completion).
    pub total_tokens: usize,

    /// Breakdown of tokens used in a completion.
    pub completion_tokens_details: Option<CompletionTokensDetails>,

    /// Breakdown of tokens used in the prompt.
    pub prompt_tokens_details: Option<PromptTokensDetails>,
}

impl FromStr for Completion {
    type Err = OapiError;

    fn from_str(content: &str) -> Result<Self, Self::Err> {
        let parse_result: Result<Self, _> = serde_json::from_str(content)
            .map_err(|e| OapiError::DeserializationError(e.to_string()));
        parse_result
    }
}
